import { baseProcedure, createTRPCRouter } from "../trpc/init";
import { createClient } from "../../../sdk/client";
import { createTask } from "../../../sdk/task";
import text2img from '@assets/text2img.json';
import text2imgHyper from '@assets/text2img_hyper.json';
import img2img from '@assets/img2img.json';
import sharp from 'sharp';
import { z } from 'zod';
import Fuse from "fuse.js";
import tags from "@assets/tags.json";
import { ModelConfigSchema } from "../../../types/schema";
import { llmOnceCall } from "../../../utils/llm";


export interface Tag {
  content: string;
  groupId: number;
  count: number;
  aliases: string;
  translation: string;
}

const fuse = new Fuse<Tag>(tags as any,  {
  keys: ["content", "aliases", "translation"],
  threshold: 0.1,
  includeScore: true
});


async function generateImage(workflow: any, options: Record<string, any>) {
  const client = createClient();
  const task = createTask(client, workflow, options);
  return task.start({ resultType: 'images' }).finally(() => { 
    client.wsInstance?.close() 
  });
}

const Text2ImgInput = z.object({
  promptText: z.string(),
  width: z.number(),
  height: z.number(),
  model: z.string(),
  isHyper: z.boolean()
});

const Img2ImgInput = z.object({
  model: z.string(),
  image: z.string(),
  similarity: z.number(),
  promptText: z.string()
});

export const imageRouter = createTRPCRouter({
  generateImage: baseProcedure
    .input(z.object({
      prompt: z.any(),
      options: z.record(z.any(), z.any())
    }))
    .mutation(async ({ input }) => {
      return generateImage(input.prompt, input.options);
    }),

  text2img: baseProcedure
    .input(Text2ImgInput)
    .mutation(async ({ input }) => {
      const workflow = input.isHyper ? text2imgHyper : text2img;
      return generateImage(workflow, {
        positive: {
          text: 'score_9, score_8_up, score_7_up, score_6_up, masterpiece, best quality, newest, absurdres, highres,' + input.promptText,
        },
        size: {
          width: input.width,
          height: input.height,
        },
        model: {
          ckpt_name: input.model,
        }
      });
    }),

  img2img: baseProcedure
    .input(Img2ImgInput)
    .mutation(async ({ input }) => {
      return generateImage(img2img, {
        model: {
          ckpt_name: input.model,
        },
        input_image: {
          image: input.image,
        },
        KSampler: {
          denoise: input.similarity/100,
        },
        positive: {
          text_c: input.promptText,
        }
      });
    }),

  getModels: baseProcedure
    .query(async () => {
      const client = createClient({ disableWs: true });
      return client.getCheckPoints();
    }),

  uploadImage: baseProcedure
    .input(z.instanceof(FormData))
    .mutation(async ({ input }) => {
      const client = createClient({ disableWs: true });
      const imageFile = input.get('image') as File;
      const arrayBuffer = await imageFile.arrayBuffer();
      const buffer = Buffer.from(arrayBuffer);

      const metadata = await sharp(buffer).metadata();
      let processedBuffer: Buffer = buffer;

      if (metadata.width && metadata.height && 
          (metadata.width > 1500 || metadata.height > 1500)) {
        const scale = Math.min(1500 / metadata.width, 1500 / metadata.height);
        processedBuffer = await sharp(buffer)
          .resize({
            width: Math.round(metadata.width * scale),
            height: Math.round(metadata.height * scale)
          })
          .toBuffer();
      }

      const newFormData = new FormData();
      newFormData.append('image', new Blob([Buffer.from(processedBuffer)]), imageFile.name);
      return client.uploadImage(newFormData);
    }),

  getQueueCount: baseProcedure
    .query(async () => {
      const client = createClient({ disableWs: true });
      const res = await client.getQueue();
      return res.queuePending.length + res.queueRunning.length;
    }),
  
  searchTags: baseProcedure
    .input(z.object({
      query: z.string()
    }))
    .query(async ({ input }) => {
      if (!input.query.trim()) return [];

      const results = fuse.search(input.query);
      return results
        .map(item => item.item)
        .sort((a, b) => b.count - a.count)
        .slice(0, 20);
    }),
  generatePrompt: baseProcedure
    .input(z.object({
      description: z.string().min(1, "描述不能为空"),
      modelConfig: ModelConfigSchema
    }))
    .mutation(async ({ input }) => {
      const { description, modelConfig } = input;

      return llmOnceCall(modelConfig, [{
        role: 'system',
        content: `你是一个Stable Diffusion提示词生成器。根据用户描述生成适合的英文提示词，格式为逗号分隔的关键词列表。直接给我提示词，不要有别的废话`
      }, {
        role: 'user',
        content: description
      }]);

    })
});